import cv2

assert cv2.__version__[0] == '4'
import numpy as np
import os
import glob


def get_K_and_D(checkerboard, imgsPath):
    CHECKERBOARD = checkerboard
    subpix_criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.1)
    calibration_flags = cv2.fisheye.CALIB_RECOMPUTE_EXTRINSIC + cv2.fisheye.CALIB_CHECK_COND + cv2.fisheye.CALIB_FIX_SKEW
    objp = np.zeros((1, CHECKERBOARD[0] * CHECKERBOARD[1], 3), np.float32)
    objp[0, :, :2] = np.mgrid[0:CHECKERBOARD[0], 0:CHECKERBOARD[1]].T.reshape(-1, 2)
    _img_shape = None
    objpoints = []
    imgpoints = []
    images = glob.glob(imgsPath + '/*.jpg')
    
    output_dir = 'images_processing/img_out/undistorted/'
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    for fname in images:
        img = cv2.imread(fname)
        if _img_shape is None:
            _img_shape = img.shape[:2]
        else:
            assert _img_shape == img.shape[:2], "All images must share the same size."

        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        ret, corners = cv2.findChessboardCorners(gray, CHECKERBOARD,
                                                 cv2.CALIB_CB_ADAPTIVE_THRESH + cv2.CALIB_CB_FAST_CHECK + cv2.CALIB_CB_NORMALIZE_IMAGE)
        if ret:
            objpoints.append(objp)
            cv2.cornerSubPix(gray, corners, (3, 3), (-1, -1), subpix_criteria)
            imgpoints.append(corners)

            # Draw and display the corners
            cv2.drawChessboardCorners(img, (CHECKERBOARD[0], CHECKERBOARD[1]), corners, ret)
            output_path = os.path.join(output_dir, os.path.basename(fname))
            cv2.imwrite(output_path, img)

            # Optional: Display the image with corners
            # cv2.imshow('findCorners', img)
            # if cv2.waitKey(0) & 0xFF == ord('q'):
            #     break

    N_OK = len(objpoints)
    K = np.zeros((3, 3))
    D = np.zeros((4, 1))
    rvecs = [np.zeros((1, 1, 3), dtype=np.float64) for i in range(N_OK)]
    tvecs = [np.zeros((1, 1, 3), dtype=np.float64) for i in range(N_OK)]
    rms, _, _, _, _ = \
        cv2.fisheye.calibrate(
            objpoints,
            imgpoints,
            gray.shape[::-1],
            K,
            D,
            rvecs,
            tvecs,
            calibration_flags,
            (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 1e-6)
        )
    DIM = _img_shape[::-1]
    print("Found " + str(N_OK) + " valid images for calibration")
    print("DIM=" + str(_img_shape[::-1]))
    print("K=np.array(" + str(K.tolist()) + ")")
    print("D=np.array(" + str(D.tolist()) + ")")

    return DIM, K, D


def undistort(img_path, K, D, DIM, scale=0.3, save_path=None):
    img = cv2.imread(img_path)
    dim1 = img.shape[:2][::-1]  # dim1 is the dimension of input image to un-distort
    assert dim1[0] / dim1[1] == DIM[0] / DIM[
        1], "Image to undistort needs to have same aspect ratio as the ones used in calibration"
    if dim1[0] != DIM[0]:
        img = cv2.resize(img, DIM, interpolation=cv2.INTER_AREA)
    Knew = K.copy()
    if scale:  # change fov
        Knew[(0, 1), (0, 1)] = scale * Knew[(0, 1), (0, 1)]
    map1, map2 = cv2.fisheye.initUndistortRectifyMap(K, D, np.eye(3), Knew, DIM, cv2.CV_16SC2)
    undistorted_img = cv2.remap(img, map1, map2, interpolation=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT)
    if save_path:
        cv2.imwrite(save_path, undistorted_img)
    else:
        # 如果未提供保存路径，则可以选择显示图像（这里已移除原imshow调用）
        # cv2.imshow("undistorted", undistorted_img)
        # cv2.waitKey(0)  # 等待按键事件，以便在显示图像后能够关闭窗口
        # cv2.destroyAllWindows()  # 关闭所有OpenCV创建的窗口
        pass  # 不执行任何操作
    return undistorted_img


# 计算内参和矫正系数
'''
# checkerboard  棋盘格的格点数目
# imgsPath: 存放鱼眼图片的路径
'''
dim, k, d = get_K_and_D((6, 9), 'images_processing/img_in/zzy_img/test')

# 去畸变
undistort('images_processing/img_in/zzy_img/test/9.jpg', k, d, dim, 1, "images_processing/img_out/undistorted_img.jpg")

if cv2.waitKey(0) & 0xFF == ord('q'):
    exit()
